

<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN"
  "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">


<html xmlns="http://www.w3.org/1999/xhtml">
  <head>
    <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
    
    <title>AG_fft_tools.convolve_nd &mdash; agpy 0.1.2 documentation</title>
    
    <link rel="stylesheet" href="../../_static/extra.css" type="text/css" />
    <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
    
    <script type="text/javascript">
      var DOCUMENTATION_OPTIONS = {
        URL_ROOT:    '../../',
        VERSION:     '0.1.2',
        COLLAPSE_INDEX: false,
        FILE_SUFFIX: '.html',
        HAS_SOURCE:  true
      };
    </script>
    <script type="text/javascript" src="../../_static/jquery.js"></script>
    <script type="text/javascript" src="../../_static/underscore.js"></script>
    <script type="text/javascript" src="../../_static/doctools.js"></script>
    <link rel="top" title="agpy 0.1.2 documentation" href="../../index.html" />
    <link rel="up" title="Module code" href="../index.html" />
     
    <script type="text/javascript">

      var _gaq = _gaq || [];
      _gaq.push(['_setDomainName', 'pyspeckit.bitbucket.org']);
      _gaq.push(['_setAllowHash', false]);
      _gaq.push(['_trackPageview']);


    </script>
    <link rel="stylesheet" type="text/css" href="../../_static/extra.css" />
  </head>
  <body>
    <div class="header-wrapper">
      <div class="header">
        <h1><a href="../../index.html">agpy 0.1.2 documentation</a></h1>
        <div class="rel">
          <a href="http://agpy.googlecode.com">agpy Home </a>  |
          <a href=../../index.html>Docs Home </a>  |
          <a href="http://code.google.com/p/agpy/w/list">Wiki</a>  |
          <a href=../../search.html>Search </a>
        </div>
       </div>
    </div>

    <div class="content-wrapper">
      <div class="content">
        <div class="document">
            
      <div class="documentwrapper">
        <div class="bodywrapper">
          <div class="body">
            
  <h1>Source code for AG_fft_tools.convolve_nd</h1><div class="highlight"><pre>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="kn">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">warnings</span>

<span class="k">try</span><span class="p">:</span>
    <span class="kn">import</span> <span class="nn">fftw3</span>
    <span class="n">has_fftw</span> <span class="o">=</span> <span class="bp">True</span>

<div class="viewcode-block" id="fftwn"><a class="viewcode-back" href="../../fft_tools.html#AG_fft_tools.convolve_nd.fftwn">[docs]</a>    <span class="k">def</span> <span class="nf">fftwn</span><span class="p">(</span><span class="n">array</span><span class="p">,</span> <span class="n">nthreads</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
        <span class="n">array</span> <span class="o">=</span> <span class="n">array</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s">&#39;complex&#39;</span><span class="p">)</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
        <span class="n">outarray</span> <span class="o">=</span> <span class="n">array</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
        <span class="n">fft_forward</span> <span class="o">=</span> <span class="n">fftw3</span><span class="o">.</span><span class="n">Plan</span><span class="p">(</span><span class="n">array</span><span class="p">,</span> <span class="n">outarray</span><span class="p">,</span> <span class="n">direction</span><span class="o">=</span><span class="s">&#39;forward&#39;</span><span class="p">,</span>
                <span class="n">flags</span><span class="o">=</span><span class="p">[</span><span class="s">&#39;estimate&#39;</span><span class="p">],</span> <span class="n">nthreads</span><span class="o">=</span><span class="n">nthreads</span><span class="p">)</span>
        <span class="n">fft_forward</span><span class="o">.</span><span class="n">execute</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">outarray</span>
</div>
<div class="viewcode-block" id="ifftwn"><a class="viewcode-back" href="../../fft_tools.html#AG_fft_tools.convolve_nd.ifftwn">[docs]</a>    <span class="k">def</span> <span class="nf">ifftwn</span><span class="p">(</span><span class="n">array</span><span class="p">,</span> <span class="n">nthreads</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
        <span class="n">array</span> <span class="o">=</span> <span class="n">array</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s">&#39;complex&#39;</span><span class="p">)</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
        <span class="n">outarray</span> <span class="o">=</span> <span class="n">array</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
        <span class="n">fft_backward</span> <span class="o">=</span> <span class="n">fftw3</span><span class="o">.</span><span class="n">Plan</span><span class="p">(</span><span class="n">array</span><span class="p">,</span> <span class="n">outarray</span><span class="p">,</span> <span class="n">direction</span><span class="o">=</span><span class="s">&#39;backward&#39;</span><span class="p">,</span>
                <span class="n">flags</span><span class="o">=</span><span class="p">[</span><span class="s">&#39;estimate&#39;</span><span class="p">],</span> <span class="n">nthreads</span><span class="o">=</span><span class="n">nthreads</span><span class="p">)</span>
        <span class="n">fft_backward</span><span class="o">.</span><span class="n">execute</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">outarray</span> <span class="o">/</span> <span class="n">np</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">array</span><span class="p">)</span></div>
<span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
    <span class="n">fftn</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">fftn</span>
    <span class="n">ifftn</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">ifftn</span>
    <span class="n">has_fftw</span> <span class="o">=</span> <span class="bp">False</span>
<span class="c"># I performed some fft speed tests and found that scipy is slower than numpy</span>
<span class="c"># http://code.google.com/p/agpy/source/browse/trunk/tests/test_ffts.py However,</span>
<span class="c"># the speed varied on machines - YMMV.  If someone finds that scipy&#39;s fft is</span>
<span class="c"># faster, we should add that as an option here... not sure how exactly</span>


<div class="viewcode-block" id="convolvend"><a class="viewcode-back" href="../../fft_tools.html#AG_fft_tools.convolve_nd.convolvend">[docs]</a><span class="k">def</span> <span class="nf">convolvend</span><span class="p">(</span><span class="n">array</span><span class="p">,</span> <span class="n">kernel</span><span class="p">,</span> <span class="n">crop</span><span class="o">=</span><span class="bp">True</span><span class="p">,</span> <span class="n">return_fft</span><span class="o">=</span><span class="bp">False</span><span class="p">,</span> <span class="n">fftshift</span><span class="o">=</span><span class="bp">True</span><span class="p">,</span>
        <span class="n">fft_pad</span><span class="o">=</span><span class="bp">True</span><span class="p">,</span> <span class="n">psf_pad</span><span class="o">=</span><span class="bp">False</span><span class="p">,</span> <span class="n">ignore_nan</span><span class="o">=</span><span class="bp">False</span><span class="p">,</span> <span class="n">quiet</span><span class="o">=</span><span class="bp">False</span><span class="p">,</span>
        <span class="n">ignore_edge_zeros</span><span class="o">=</span><span class="bp">True</span><span class="p">,</span> <span class="n">min_wt</span><span class="o">=</span><span class="mf">1e-8</span><span class="p">,</span> <span class="n">force_ignore_zeros_off</span><span class="o">=</span><span class="bp">False</span><span class="p">,</span>
        <span class="n">normalize_kernel</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">,</span> <span class="n">use_numpy_fft</span><span class="o">=</span><span class="ow">not</span> <span class="n">has_fftw</span><span class="p">,</span> <span class="n">nthreads</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Convolve an ndarray with an nd-kernel.  Returns a convolved image with shape =</span>
<span class="sd">    array.shape.  Assumes image &amp; kernel are centered.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    array: `numpy.ndarray`</span>
<span class="sd">        Array to be convolved with *kernel*</span>
<span class="sd">    kernel: `numpy.ndarray`</span>
<span class="sd">        Will be normalized if *normalize_kernel* is set.  Assumed to be</span>
<span class="sd">        centered (i.e., shifts may result if your kernel is asymmetric)</span>

<span class="sd">    Options</span>
<span class="sd">    -------</span>
<span class="sd">    fft_pad: bool</span>
<span class="sd">        Default on.  Zero-pad image to the nearest 2^n</span>
<span class="sd">    psf_pad: bool</span>
<span class="sd">        Default off.  Zero-pad image to be at least the sum of the image sizes</span>
<span class="sd">        (in order to avoid edge-wrapping when smoothing)</span>
<span class="sd">    crop: bool</span>
<span class="sd">        Default on.  Return an image of the size of the largest input image.</span>
<span class="sd">        If the images are asymmetric in opposite directions, will return the</span>
<span class="sd">        largest image in both directions.</span>
<span class="sd">        For example, if an input image has shape [100,3] but a kernel with shape</span>
<span class="sd">        [6,6] is used, the output will be [100,6].</span>
<span class="sd">    return_fft: bool</span>
<span class="sd">        Return the fft(image)*fft(kernel) instead of the convolution (which is</span>
<span class="sd">        ifft(fft(image)*fft(kernel))).  Useful for making PSDs.</span>
<span class="sd">    fftshift: bool</span>
<span class="sd">        If return_fft on, will shift &amp; crop image to appropriate dimensions</span>
<span class="sd">    ignore_nan: bool</span>
<span class="sd">        attempts to re-weight assuming NAN values are meant to be ignored, not</span>
<span class="sd">        treated as zero.  If this is off, all NaN values will be treated as</span>
<span class="sd">        zero.</span>
<span class="sd">    ignore_edge_zeros: bool</span>
<span class="sd">        Ignore the zero-pad-created zeros.  This will effectively decrease</span>
<span class="sd">        the kernel area on the edges but will not re-normalize the kernel.</span>
<span class="sd">        This is on by default but I&#39;m not entirely sure it should be...</span>
<span class="sd">    force_ignore_zeros_off: bool</span>
<span class="sd">        You can choose to turn off the ignore-zeros when padding; this may be</span>
<span class="sd">        desirable if you want to think of the region outside of your image as</span>
<span class="sd">        all zeros</span>
<span class="sd">    min_wt: float</span>
<span class="sd">        If ignoring nans/zeros, force all grid points with a weight less than</span>
<span class="sd">        this value to NAN (the weight of a grid point with *no* ignored</span>
<span class="sd">        neighbors is 1.0).</span>
<span class="sd">    normalize_kernel: function</span>
<span class="sd">        if specified, function to divide kernel by to normalize it.  e.g.,</span>
<span class="sd">        normalize_kernel=np.sum means that kernel will be modified to be:</span>
<span class="sd">        kernel = kernel / np.sum(kernel)</span>
<span class="sd">    nthreads: int</span>
<span class="sd">        if fftw3 is installed, can specify the number of threads to allow FFTs</span>
<span class="sd">        to use.  Probably only helpful for large arrays</span>
<span class="sd">    use_numpy_fft: bool</span>
<span class="sd">        Force the code to use the numpy FFTs instead of FFTW even if FFTW is</span>
<span class="sd">        installed</span>

<span class="sd">    Returns</span>
<span class="sd">    -------</span>
<span class="sd">    default: *array* convolved with *kernel*</span>
<span class="sd">    if return_fft: fft(*array*) * fft(*kernel*)</span>
<span class="sd">      * if fftshift: Determines whether the fft will be shifted before</span>
<span class="sd">        returning</span>
<span class="sd">    if *crop* == False : Returns the image, but with the fft-padded size</span>
<span class="sd">      instead of the input size</span>

<span class="sd">    .. note: This is also implemented in :astropy:nddata:</span>

<span class="sd">    &#39;&#39;&#39;</span>

    <span class="c"># Checking copied from convolve.py - however, since FFTs have real &amp;</span>
    <span class="c"># complex components, we change the types.  Only the real part will be</span>
    <span class="c"># returned!</span>
    <span class="c"># Check that the arguments are lists or Numpy arrays</span>
    <span class="n">array</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">array</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">complex</span><span class="p">)</span>
    <span class="n">kernel</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">kernel</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">complex</span><span class="p">)</span>

    <span class="c"># Check that the number of dimensions is compatible</span>
    <span class="k">if</span> <span class="n">array</span><span class="o">.</span><span class="n">ndim</span> <span class="o">!=</span> <span class="n">kernel</span><span class="o">.</span><span class="n">ndim</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s">&#39;array and kernel have differing number of&#39;</span>
                        <span class="s">&#39;dimensions&#39;</span><span class="p">)</span>

    <span class="c"># store the dtype for conversion back later</span>
    <span class="n">array_dtype</span> <span class="o">=</span> <span class="n">array</span><span class="o">.</span><span class="n">dtype</span>
    <span class="c"># turn the arrays into &#39;complex&#39; arrays</span>
    <span class="k">if</span> <span class="n">array</span><span class="o">.</span><span class="n">dtype</span><span class="o">.</span><span class="n">kind</span> <span class="o">!=</span> <span class="s">&#39;c&#39;</span><span class="p">:</span>
        <span class="n">array</span> <span class="o">=</span> <span class="n">array</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">complex</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">kernel</span><span class="o">.</span><span class="n">dtype</span><span class="o">.</span><span class="n">kind</span> <span class="o">!=</span> <span class="s">&#39;c&#39;</span><span class="p">:</span>
        <span class="n">kernel</span> <span class="o">=</span> <span class="n">kernel</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">complex</span><span class="p">)</span>

    <span class="c"># mask catching - masks must be turned into NaNs for use later</span>
    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">ma</span><span class="o">.</span><span class="n">is_masked</span><span class="p">(</span><span class="n">array</span><span class="p">):</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">array</span><span class="o">.</span><span class="n">mask</span>
        <span class="n">array</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">array</span><span class="p">)</span>
        <span class="n">array</span><span class="p">[</span><span class="n">mask</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">ma</span><span class="o">.</span><span class="n">is_masked</span><span class="p">(</span><span class="n">kernel</span><span class="p">):</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">kernel</span><span class="o">.</span><span class="n">mask</span>
        <span class="n">kernel</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span>
        <span class="n">kernel</span><span class="p">[</span><span class="n">mask</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>

    <span class="c"># replace fftn if has_fftw so that nthreads can be passed</span>
    <span class="k">global</span> <span class="n">fftn</span><span class="p">,</span> <span class="n">ifftn</span>
    <span class="k">if</span> <span class="n">has_fftw</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">use_numpy_fft</span><span class="p">:</span>

        <span class="k">def</span> <span class="nf">fftn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
            <span class="k">return</span> <span class="n">fftwn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="n">nthreads</span><span class="o">=</span><span class="n">nthreads</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

        <span class="k">def</span> <span class="nf">ifftn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
            <span class="k">return</span> <span class="n">ifftwn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="n">nthreads</span><span class="o">=</span><span class="n">nthreads</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">use_numpy_fft</span><span class="p">:</span>
        <span class="n">fftn</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">fftn</span>
        <span class="n">ifftn</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">ifftn</span>


    <span class="c"># NAN catching</span>
    <span class="n">nanmaskarray</span> <span class="o">=</span> <span class="p">(</span><span class="n">array</span> <span class="o">!=</span> <span class="n">array</span><span class="p">)</span>
    <span class="n">array</span><span class="p">[</span><span class="n">nanmaskarray</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="n">nanmaskkernel</span> <span class="o">=</span> <span class="p">(</span><span class="n">kernel</span> <span class="o">!=</span> <span class="n">kernel</span><span class="p">)</span>
    <span class="n">kernel</span><span class="p">[</span><span class="n">nanmaskkernel</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">if</span> <span class="p">((</span><span class="n">nanmaskarray</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">nanmaskkernel</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">ignore_nan</span>
            <span class="ow">and</span> <span class="ow">not</span> <span class="n">quiet</span><span class="p">):</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s">&quot;NOT ignoring nan values even though they are present&quot;</span> <span class="o">+</span>
                <span class="s">&quot; (they are treated as 0)&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="p">((</span><span class="n">psf_pad</span> <span class="ow">or</span> <span class="n">fft_pad</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">ignore_edge_zeros</span> <span class="ow">and</span> <span class="ow">not</span>
            <span class="n">force_ignore_zeros_off</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">quiet</span><span class="p">):</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s">&quot;when psf_pad or fft_pad are enabled, &quot;</span> <span class="o">+</span>
                <span class="s">&quot;ignore_edge_zeros is forced on&quot;</span><span class="p">)</span>
        <span class="n">ignore_edge_zeros</span> <span class="o">=</span> <span class="bp">True</span>
    <span class="k">elif</span> <span class="n">force_ignore_zeros_off</span><span class="p">:</span>
        <span class="n">ignore_edge_zeros</span> <span class="o">=</span> <span class="bp">False</span>

    <span class="k">if</span> <span class="n">normalize_kernel</span><span class="p">:</span>
        <span class="c"># try this.  If a function is not passed, the code will just crash... I</span>
        <span class="c"># think type checking would be better but PEPs say otherwise...</span>
        <span class="n">kernel</span> <span class="o">=</span> <span class="n">kernel</span> <span class="o">/</span> <span class="n">normalize_kernel</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span>


    <span class="n">arrayshape</span> <span class="o">=</span> <span class="n">array</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">kernshape</span> <span class="o">=</span> <span class="n">kernel</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">ndim</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">array</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">ndim</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">kernshape</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s">&quot;Image and kernel must &quot;</span> <span class="o">+</span>
            <span class="s">&quot;have same number of dimensions&quot;</span><span class="p">)</span>
    <span class="c"># find ideal size (power of 2) for fft.</span>
    <span class="c"># Can add shapes because they are tuples</span>
    <span class="k">if</span> <span class="n">fft_pad</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">psf_pad</span><span class="p">:</span>
            <span class="c"># add the dimensions and then take the max (bigger)</span>
            <span class="n">fsize</span> <span class="o">=</span> <span class="mi">2</span><span class="o">**</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">log2</span><span class="p">(</span>
                <span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">arrayshape</span><span class="p">)</span> <span class="o">+</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">kernshape</span><span class="p">))))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c"># add the shape lists (max of a list of length 4) (smaller)</span>
            <span class="c"># also makes the shapes square</span>
            <span class="n">fsize</span> <span class="o">=</span> <span class="mi">2</span><span class="o">**</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">log2</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">arrayshape</span><span class="o">+</span><span class="n">kernshape</span><span class="p">)))</span>
        <span class="n">newshape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="n">fsize</span> <span class="k">for</span> <span class="n">ii</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">ndim</span><span class="p">)])</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">psf_pad</span><span class="p">:</span>
            <span class="c"># just add the biggest dimensions</span>
            <span class="n">newshape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">arrayshape</span><span class="p">)</span><span class="o">+</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">kernshape</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">newshape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">([</span><span class="n">imsh</span><span class="p">,</span> <span class="n">kernsh</span><span class="p">])</span>
                <span class="k">for</span> <span class="n">imsh</span><span class="p">,</span> <span class="n">kernsh</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">arrayshape</span><span class="p">,</span> <span class="n">kernshape</span><span class="p">)])</span>


    <span class="c"># separate each dimension by the padding size...  this is to determine the</span>
    <span class="c"># appropriate slice size to get back to the input dimensions</span>
    <span class="n">arrayslices</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="n">kernslices</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">ii</span><span class="p">,</span> <span class="p">(</span><span class="n">newdimsize</span><span class="p">,</span> <span class="n">arraydimsize</span><span class="p">,</span> <span class="n">kerndimsize</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">newshape</span><span class="p">,</span> <span class="n">arrayshape</span><span class="p">,</span> <span class="n">kernshape</span><span class="p">)):</span>
        <span class="n">center</span> <span class="o">=</span> <span class="n">newdimsize</span><span class="o">/</span><span class="mf">2.</span>
        <span class="n">arrayslices</span> <span class="o">+=</span> <span class="p">[</span><span class="nb">slice</span><span class="p">(</span><span class="n">center</span> <span class="o">-</span> <span class="n">arraydimsize</span><span class="o">/</span><span class="mf">2.</span><span class="p">,</span>
            <span class="n">center</span> <span class="o">+</span> <span class="n">arraydimsize</span><span class="o">/</span><span class="mf">2.</span><span class="p">)]</span>
        <span class="n">kernslices</span> <span class="o">+=</span> <span class="p">[</span><span class="nb">slice</span><span class="p">(</span><span class="n">center</span> <span class="o">-</span> <span class="n">kerndimsize</span><span class="o">/</span><span class="mf">2.</span><span class="p">,</span>
            <span class="n">center</span> <span class="o">+</span> <span class="n">kerndimsize</span><span class="o">/</span><span class="mf">2.</span><span class="p">)]</span>

    <span class="n">bigarray</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">newshape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">complex128</span><span class="p">)</span>
    <span class="n">bigkernel</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">newshape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">complex128</span><span class="p">)</span>
    <span class="n">bigarray</span><span class="p">[</span><span class="n">arrayslices</span><span class="p">]</span> <span class="o">=</span> <span class="n">array</span>
    <span class="n">bigkernel</span><span class="p">[</span><span class="n">kernslices</span><span class="p">]</span> <span class="o">=</span> <span class="n">kernel</span>
    <span class="n">arrayfft</span> <span class="o">=</span> <span class="n">fftn</span><span class="p">(</span><span class="n">bigarray</span><span class="p">)</span>
    <span class="n">kernfft</span> <span class="o">=</span> <span class="n">fftn</span><span class="p">(</span><span class="n">bigkernel</span><span class="p">)</span>
    <span class="n">fftmult</span> <span class="o">=</span> <span class="n">arrayfft</span><span class="o">*</span><span class="n">kernfft</span>
    <span class="k">if</span> <span class="n">ignore_nan</span> <span class="ow">or</span> <span class="n">ignore_edge_zeros</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">ignore_edge_zeros</span><span class="p">:</span>
            <span class="n">bigimwt</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">newshape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">complex128</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">bigimwt</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">newshape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">complex128</span><span class="p">)</span>
        <span class="n">bigimwt</span><span class="p">[</span><span class="n">arrayslices</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="o">-</span><span class="n">nanmaskarray</span><span class="o">*</span><span class="n">ignore_nan</span>
        <span class="n">wtfft</span> <span class="o">=</span> <span class="n">fftn</span><span class="p">(</span><span class="n">bigimwt</span><span class="p">)</span>
        <span class="c"># I think this one HAS to be normalized (i.e., the weights can&#39;t be</span>
        <span class="c"># computed with a non-normalized kernel)</span>
        <span class="n">wtfftmult</span> <span class="o">=</span> <span class="n">wtfft</span><span class="o">*</span><span class="n">kernfft</span><span class="o">/</span><span class="n">kernel</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
        <span class="n">wtsm</span> <span class="o">=</span> <span class="n">ifftn</span><span class="p">(</span><span class="n">wtfftmult</span><span class="p">)</span>
        <span class="c"># need to re-zero weights outside of the image (if it is padded, we</span>
        <span class="c"># still don&#39;t weight those regions)</span>
        <span class="n">bigimwt</span><span class="p">[</span><span class="n">arrayslices</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">fftshift</span><span class="p">(</span><span class="n">wtsm</span><span class="p">)</span><span class="o">.</span><span class="n">real</span><span class="p">[</span><span class="n">arrayslices</span><span class="p">]</span>

    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">fftmult</span><span class="p">)</span><span class="o">.</span><span class="n">any</span><span class="p">():</span>
        <span class="c"># this check should be unnecessary; call it an insanity check</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s">&quot;Encountered NaNs in convolve.  This is disallowed.&quot;</span><span class="p">)</span>

    <span class="c"># restore nans in original image (they were modified inplace earlier)</span>
    <span class="c"># We don&#39;t have to worry about masked arrays - if input was masked, it was</span>
    <span class="c"># copied</span>
    <span class="n">array</span><span class="p">[</span><span class="n">nanmaskarray</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
    <span class="n">kernel</span><span class="p">[</span><span class="n">nanmaskkernel</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>

    <span class="k">if</span> <span class="n">return_fft</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">fftshift</span><span class="p">:</span> <span class="c"># default on</span>
            <span class="k">if</span> <span class="n">crop</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">fftshift</span><span class="p">(</span><span class="n">fftmult</span><span class="p">)[</span><span class="n">arrayslices</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">fftshift</span><span class="p">(</span><span class="n">fftmult</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">fftmult</span>

    <span class="k">if</span> <span class="n">ignore_nan</span> <span class="ow">or</span> <span class="n">ignore_edge_zeros</span><span class="p">:</span>
        <span class="n">rifft</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">fftshift</span><span class="p">(</span><span class="n">ifftn</span><span class="p">(</span><span class="n">fftmult</span><span class="p">))</span> <span class="o">/</span> <span class="n">bigimwt</span>
        <span class="n">rifft</span><span class="p">[</span><span class="n">bigimwt</span> <span class="o">&lt;</span> <span class="n">min_wt</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">rifft</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fft</span><span class="o">.</span><span class="n">fftshift</span><span class="p">(</span><span class="n">ifftn</span><span class="p">(</span><span class="n">fftmult</span><span class="p">))</span>

    <span class="k">if</span> <span class="n">crop</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">rifft</span><span class="p">[</span><span class="n">arrayslices</span><span class="p">]</span><span class="o">.</span><span class="n">real</span>
        <span class="k">return</span> <span class="n">result</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">rifft</span><span class="o">.</span><span class="n">real</span>



</div>
<span class="kn">import</span> <span class="nn">pytest</span>
<span class="kn">import</span> <span class="nn">itertools</span>
<span class="n">params</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">itertools</span><span class="o">.</span><span class="n">product</span><span class="p">((</span><span class="bp">True</span><span class="p">,</span><span class="bp">False</span><span class="p">),(</span><span class="bp">True</span><span class="p">,</span><span class="bp">False</span><span class="p">),(</span><span class="bp">True</span><span class="p">,</span><span class="bp">False</span><span class="p">)))</span>
<span class="nd">@pytest.mark.parametrize</span><span class="p">((</span><span class="s">&#39;psf_pad&#39;</span><span class="p">,</span><span class="s">&#39;use_numpy_fft&#39;</span><span class="p">,</span><span class="s">&#39;force_ignore_zeros_off&#39;</span><span class="p">),</span><span class="n">params</span><span class="p">)</span>
<div class="viewcode-block" id="test_3d"><a class="viewcode-back" href="../../fft_tools.html#AG_fft_tools.convolve_nd.test_3d">[docs]</a><span class="k">def</span> <span class="nf">test_3d</span><span class="p">(</span><span class="n">psf_pad</span><span class="p">,</span> <span class="n">use_numpy_fft</span><span class="p">,</span> <span class="n">force_ignore_zeros_off</span><span class="p">,</span> <span class="n">debug</span><span class="o">=</span><span class="bp">False</span><span class="p">,</span> <span class="n">tolerance</span><span class="o">=</span><span class="mf">1e-17</span><span class="p">):</span>
    <span class="n">array</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">([</span><span class="mi">32</span><span class="p">,</span><span class="mi">32</span><span class="p">,</span><span class="mi">32</span><span class="p">])</span>
    <span class="n">array</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">]</span><span class="o">=</span><span class="mi">1</span>
    <span class="n">array</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">15</span><span class="p">]</span><span class="o">=</span><span class="mi">1</span>
    <span class="n">kern</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">([</span><span class="mi">32</span><span class="p">,</span><span class="mi">32</span><span class="p">,</span><span class="mi">32</span><span class="p">])</span>
    <span class="n">kern</span><span class="p">[</span><span class="mi">14</span><span class="p">:</span><span class="mi">19</span><span class="p">,</span><span class="mi">14</span><span class="p">:</span><span class="mi">19</span><span class="p">,</span><span class="mi">14</span><span class="p">:</span><span class="mi">19</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>

    <span class="n">conv1</span> <span class="o">=</span> <span class="n">convolvend</span><span class="p">(</span><span class="n">array</span><span class="p">,</span> <span class="n">kern</span><span class="p">,</span> <span class="n">psf_pad</span><span class="o">=</span><span class="n">psf_pad</span><span class="p">,</span> <span class="n">force_ignore_zeros_off</span><span class="o">=</span><span class="n">force_ignore_zeros_off</span><span class="p">,</span> <span class="n">debug</span><span class="o">=</span><span class="n">debug</span><span class="p">)</span>

    <span class="k">print</span> <span class="s">&quot;psf_pad=</span><span class="si">%s</span><span class="s"> use_numpy=</span><span class="si">%s</span><span class="s"> force_ignore_zeros_off=</span><span class="si">%s</span><span class="s">&quot;</span> <span class="o">%</span> <span class="p">(</span><span class="n">psf_pad</span><span class="p">,</span> <span class="n">use_numpy_fft</span><span class="p">,</span> <span class="n">force_ignore_zeros_off</span><span class="p">)</span>
    <span class="k">print</span> <span class="s">&quot;side,center: </span><span class="si">%g</span><span class="s">,</span><span class="si">%g</span><span class="s">&quot;</span> <span class="o">%</span> <span class="p">(</span><span class="n">conv1</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">15</span><span class="p">],</span><span class="n">conv1</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">])</span>
    <span class="k">if</span> <span class="n">force_ignore_zeros_off</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">psf_pad</span><span class="p">:</span>
        <span class="k">assert</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">conv1</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">15</span><span class="p">]</span> <span class="o">-</span> <span class="mf">1.</span><span class="o">/</span><span class="mf">125.</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">tolerance</span><span class="p">)</span>
        <span class="k">assert</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">conv1</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">1</span><span class="p">,</span><span class="mi">15</span><span class="p">]</span> <span class="o">-</span> <span class="mf">1.</span><span class="o">/</span><span class="mf">125.</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">tolerance</span><span class="p">)</span>
        <span class="k">assert</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">conv1</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">]</span> <span class="o">-</span> <span class="mf">1.</span><span class="o">/</span><span class="mf">125.</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">tolerance</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">assert</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">conv1</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">15</span><span class="p">]</span> <span class="o">-</span> <span class="mf">1.</span><span class="o">/</span><span class="mf">75.</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">tolerance</span><span class="p">)</span>
        <span class="k">assert</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">conv1</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">1</span><span class="p">,</span><span class="mi">15</span><span class="p">]</span> <span class="o">-</span> <span class="mf">1.</span><span class="o">/</span><span class="mf">100.</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">tolerance</span><span class="p">)</span>
        <span class="k">assert</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">conv1</span><span class="p">[</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">]</span> <span class="o">-</span> <span class="mf">1.</span><span class="o">/</span><span class="mf">125.</span><span class="p">)</span> <span class="o">&lt;</span> <span class="n">tolerance</span><span class="p">)</span>

 </div>
</pre></div>

          </div>
        </div>
      </div>
        </div>
        <div class="sidebar">
          <h3>Table Of Contents</h3>
          <ul>
<li class="toctree-l1"><a class="reference internal" href="../../agpy.html">Adam Ginsburg&#8217;s Python Code (agpy)</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../image_tools.html">Image Tools</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../fft_tools.html">AG_fft_tools Package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../plfit.html">plfit Package</a></li>
</ul>

          <h3 style="margin-top: 1.5em;">Search</h3>
          <form class="search" action="../../search.html" method="get">
            <input type="text" name="q" />
            <input type="submit" value="Go" />
            <input type="hidden" name="check_keywords" value="yes" />
            <input type="hidden" name="area" value="default" />
          </form>
          <p class="searchtip" style="font-size: 90%">
            Enter search terms or a module, class or function name.
          </p>
        </div>
        <div class="clearer"></div>
      </div>
    </div>

    <div class="footer">
      &copy; Copyright 2011, Adam Ginsburg.
      Created using <a href="http://sphinx.pocoo.org/">Sphinx</a> 1.2pre.
    <script type="text/javascript">

      var _gaq = _gaq || [];
      _gaq.push(['_setAccount', 'UA-6253248-2']);
      _gaq.push(['_trackPageview']);

      (function() {
        var ga = document.createElement('script'); ga.type = 'text/javascript'; ga.async = true;
        ga.src = ('https:' == document.location.protocol ? 'https://ssl' : 'http://www') + '.google-analytics.com/ga.js';
        var s = document.getElementsByTagName('script')[0]; s.parentNode.insertBefore(ga, s);
      })();

    </script>
        
    </div>
  </body>
</html>